package com.kitten.chatglmsdk.executor.aigc;

import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.databind.util.BeanUtil;
import com.kitten.chatglmsdk.IOpenAiApi;
import com.kitten.chatglmsdk.executor.Executor;
import com.kitten.chatglmsdk.executor.result.ResultHandler;
import com.kitten.chatglmsdk.model.*;
import com.kitten.chatglmsdk.session.Configuration;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeanUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;

public class GLMExecutor implements Executor, ResultHandler {
    private static final Logger log = LoggerFactory.getLogger(GLMExecutor.class);
    private final Configuration configuration;
    private final EventSource.Factory factory;
    private IOpenAiApi openAiApi;
    private OkHttpClient okHttpClient;

    public GLMExecutor(Configuration configuration) {
        this.configuration = configuration;
        this.factory = configuration.createRequestFactory();
        this.openAiApi = configuration.getOpenAiApi();
        this.okHttpClient = configuration.getOkHttpClient();
    }
    public EventSource completions(ChatCompletionRequest chatCompletionRequest, EventSourceListener eventSourceListener) throws Exception {
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v4/chat/completions")).post(RequestBody.create(MediaType.parse("application/json; charset=utf-8"), chatCompletionRequest.toString())).build();
        return this.factory.newEventSource(request, chatCompletionRequest.getIsCompatible() ? this.eventSourceListener(eventSourceListener) : eventSourceListener);
    }

    public CompletableFuture<String> completions(ChatCompletionRequest chatCompletionRequest) throws InterruptedException {
        final CompletableFuture<String> future = new CompletableFuture();
        final StringBuffer dataBuffer = new StringBuffer();
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v4/chat/completions")).post(RequestBody.create(MediaType.parse("application/json;charset=utf-8"), chatCompletionRequest.toString())).build();
        this.factory.newEventSource(request, new EventSourceListener() {
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                if ("[DONE]".equals(data)) {
                    GLMExecutor.log.info("[输出结束] Tokens {}", JSON.toJSONString(data));
                } else {
                    ChatCompletionResponse response = (ChatCompletionResponse)JSON.parseObject(data, ChatCompletionResponse.class);
                    GLMExecutor.log.info("测试结果：{}", JSON.toJSONString(response));
                    List<ChatCompletionResponse.Choice> choices = response.getChoices();
                    Iterator var7 = choices.iterator();

                    while(var7.hasNext()) {
                        ChatCompletionResponse.Choice choice = (ChatCompletionResponse.Choice)var7.next();
                        if (!"stop".equals(choice.getFinishReason())) {
                            dataBuffer.append(choice.getDelta().getContent());
                        }
                    }

                }
            }
            public void onClosed(EventSource eventSource) {
                future.complete(dataBuffer.toString());
            }

            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }
        });
        return future;
    }



    public ChatCompletionSyncResponse completionsSync(ChatCompletionRequest chatCompletionRequest) throws IOException {
        chatCompletionRequest.setStream(false);
        log.info("stream设置完成");
        ChatNewCompletionRequest newCompletionRequest = new ChatNewCompletionRequest();
        BeanUtils.copyProperties(chatCompletionRequest, newCompletionRequest);
        newCompletionRequest.setModel(chatCompletionRequest.getModel().getCode());
        List<ChatCompletionRequest.Prompt> targetList = chatCompletionRequest.getPrompt();
        ArrayList<ChatNewCompletionRequest.Prompt> prompts = new ArrayList<>();
        for (ChatCompletionRequest.Prompt prompt : targetList) {
            ChatNewCompletionRequest.Prompt newPrompt = new ChatNewCompletionRequest.Prompt();
            BeanUtils.copyProperties(prompt, newPrompt);
            prompts.add(newPrompt);
        }
        newCompletionRequest.setMessages(prompts);
        log.info("请求体:{}", newCompletionRequest.toString());
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v4/chat/completions")).post(RequestBody.create(MediaType.parse("application/json; charset=utf-8"), newCompletionRequest.toString())).build();
        OkHttpClient okHttpClient = this.configuration.getOkHttpClient();
        Response response = okHttpClient.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new RuntimeException("Request failed");
        } else {
            return (ChatCompletionSyncResponse) JSON.parseObject(response.body().string(), ChatCompletionSyncResponse.class);
        }
    }

    /**
     * 图片生成
     * @param request
     * @return
     */
    @Override
    public ImageCompletionResponse imageGeneration(ImageCompletionRequest request) {
        log.info("图像生成入参:{}", request.toString());
        ImageCompletionResponse imageCompletionResponse = this.openAiApi.getImages(request).blockingGet();
        return imageCompletionResponse;
    }


    public EventSourceListener eventSourceListener(final EventSourceListener eventSourceListener) {
        return new EventSourceListener() {
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                if (!"[DONE]".equals(data)) {
                    ChatCompletionResponse response = (ChatCompletionResponse)JSON.parseObject(data, ChatCompletionResponse.class);
                    if (response.getChoices() != null && 1 == response.getChoices().size() && "stop".equals(((ChatCompletionResponse.Choice)response.getChoices().get(0)).getFinishReason())) {
                        eventSourceListener.onEvent(eventSource, id, EventType.finish.getCode(), data);
                    } else {
                        eventSourceListener.onEvent(eventSource, id, EventType.add.getCode(), data);
                    }
                }
            }
            public void onClosed(EventSource eventSource) {
                eventSourceListener.onClosed(eventSource);
            }

            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                eventSourceListener.onFailure(eventSource, t, response);
            }
        };
    }


}
